#pragma once
#include <vector>
#include <cmath>

#include "grid_struct.hpp"

#define YXZ(i, j, k) j + i *(GirdMax[1] - GridMin[1]) + k *(GirdMax[1] - GridMin[1]) * (GirdMax[0] - GridMin[0])

template <typename V>
static bool isDivisible(V Side_Length, V stride)
{
    if (stride == 0.0)
    {
        return false;
    }
    V quotient = Side_Length / stride;
    return std::fabs(quotient - std::round(quotient)) < 1e-6; 
}

template <typename T, typename V>
int Uniform_Search_eqDis(const std::vector<T> &SearchArray, T &target, V stride[3], V ld[3], int GirdMax[3], int GridMin[3])
{
    bool x_boundary = false, y_boundary = false, z_boundary = false;
    // int scale_x = static_cast<int>(ld[0] / stride[0]), scale_y = static_cast<int>(ld[1] / stride[1]), scale_z = static_cast<int>(ld[2] / stride[2]);
    V scale_target_x = (target.x / stride[0]), scale_target_y = (target.y / stride[1]), scale_target_z = (target.z / stride[2]);

    x_boundary = isDivisible<V>(target.x, stride[0]) ? true : false;
    y_boundary = isDivisible<V>(target.y, stride[1]) ? true : false;
    z_boundary = isDivisible<V>(target.z, stride[2]) ? true : false;

    if (x_boundary && y_boundary && z_boundary)
    {
        // Assume y-x-z cycle
        int index = YXZ(static_cast<int>(scale_target_x), static_cast<int>(scale_target_y), static_cast<int>(scale_target_z));
        target.s1 = SearchArray[index].s1;
        return 0;
    }

    int x_upper, x_lower, y_upper, y_lower, z_upper, z_lower;
    if (x_boundary)
    {
        if (target.x > (ld[0] / 2))
        {
            x_upper = static_cast<int>(scale_target_x);
            x_lower = x_upper - 1;
        }
        else
        {
            x_lower = static_cast<int>(scale_target_x);
            x_upper = x_lower + 1;
        }
    }
    else
    {
        x_upper = static_cast<int>(std::ceil(scale_target_x));
        x_lower = static_cast<int>(std::floor(scale_target_x));
    }

    if (y_boundary)
    {
        if (target.y > (ld[1] / 2))
        {
            y_upper = static_cast<int>(scale_target_y);
            y_lower = y_upper - 1;
        }
        else
        {
            y_lower = static_cast<int>(scale_target_y);
            y_upper = y_lower + 1;
        }
    }
    else
    {
        y_upper = static_cast<int>(std::ceil(scale_target_y));
        y_lower = static_cast<int>(std::floor(scale_target_y));
    }

    if (z_boundary)
    {
        if (target.z > (ld[2] / 2))
        {
            z_upper = static_cast<int>(scale_target_z);
            z_lower = z_upper - 1;
        }
        else
        {
            z_lower = static_cast<int>(scale_target_z);
            z_upper = z_lower + 1;
        }
    }
    else
    {
        z_upper = static_cast<int>(std::ceil(scale_target_z));
        z_lower = static_cast<int>(std::floor(scale_target_z));
    }

    // Assume that the order of grid points follows y-x-z cycle
    {
        V c00, c01, c10, c11;
        V x_d, y_d, z_d;
        x_d = scale_target_x - static_cast<V>(x_lower);
        y_d = scale_target_y - static_cast<V>(y_lower);
        z_d = scale_target_z - static_cast<V>(z_lower);

        //(0, 0, 0) And (1, 0, 0)
        c00 = (1.0 - x_d) * SearchArray[YXZ(x_lower, y_lower, z_lower)].s1 + 
                     x_d * SearchArray[YXZ(x_upper, y_lower, z_lower)].s1;
        //(0, 0, 1) And (1, 0, 1)
        c01 = (1.0 - x_d) * SearchArray[YXZ(x_lower, y_lower, z_upper)].s1 + 
                     x_d * SearchArray[YXZ(x_upper, y_lower, z_upper)].s1;
        //(0, 1, 0) And (1, 1, 0)
        c10 = (1.0 - x_d) * SearchArray[YXZ(x_lower, y_upper, z_lower)].s1 + 
                     x_d * SearchArray[YXZ(x_upper, y_upper, z_lower)].s1;
        //(0, 1, 1) And (1, 1, 1)
        c11 = (1.0 - x_d) * SearchArray[YXZ(x_lower, y_upper, z_upper)].s1 + 
                     x_d * SearchArray[YXZ(x_upper, y_upper, z_upper)].s1;
        // c0 = (c00 * (1.0 - y_d) + c10 * y_d);
        // c1 = (c01 * (1.0 - y_d) + c11 * y_d);
        target.s1 = (c00 * (1.0 - y_d) + c10 * y_d) * (1.0 - z_d) + (c01 * (1.0 - y_d) + c11 * y_d) * z_d;
    }

    return 0;
}

template <typename T, typename V>
void Compute_prefix(const std::vector<T> &SearchArray, int numGridD1, int numGridD2, int numGridD3, V **prefixD1, V **prefixD2, V **prefixD3)
{
    int i, index;
    *prefixD1 = (V *)malloc(sizeof(V) * numGridD1);
    *prefixD2 = (V *)malloc(sizeof(V) * numGridD2);
    *prefixD3 = (V *)malloc(sizeof(V) * numGridD3);
    // Assume y-x-z cycle
    for(i = 0 ; i < numGridD2 ; i ++)
    {
        (*prefixD2)[i] = SearchArray[i].y;
    }
    for(i -= 1, index = 0 ; i < numGridD2 * numGridD1 ; i += numGridD2, index ++)
    {
        (*prefixD1)[index] = SearchArray[i].x;
    }
    for(i -= numGridD2, index = 0; i < numGridD2 * numGridD1 * numGridD3 ; i += numGridD2 * numGridD1, index ++)
    {
        (*prefixD3)[index] = SearchArray[i].z;
    }
}

template <typename V>
// LorR = 0: left, LorR = 1: right
static void BinSearch(V *prefix, V target, int &lower, int &upper, bool &on_boundary, int LorR)
{
    int left = lower, right = upper, mid;

    while (left < right)
    {
        mid = (left + right) / 2;
        if (prefix[mid] == target)
        {
            on_boundary = true;
            if (LorR == 0) { lower = mid, upper = mid + 1;} 
            else { upper = mid, lower = mid - 1;}
            return;
        }
        else if (prefix[mid] < target)
        {
            left = mid + 1;
        }
        else
        {
            right = mid;
        }
    }

    lower = left - 1;
    upper = left;
}

template <typename T, typename V>
int Uniform_Search_uneqDis(const std::vector<T> &SearchArray, T &target, V *prefixD1, V *prefixD2, V *prefixD3, V ld[3], int GirdMax[3], int GridMin[3])
{
    bool x_boundary = false, y_boundary = false, z_boundary = false;
    int x_upper = GirdMax[0], x_lower = GridMin[0], y_upper = GirdMax[1], y_lower = GridMin[1], z_upper = GirdMax[2], z_lower = GridMin[2];
    int index = 0;

    if (target.x > (ld[0] / 2))
    {
        BinSearch(prefixD1, target.x, x_lower, x_upper, x_boundary, 1);
        index += x_upper * (GirdMax[1] - GridMin[1]);
    }
    else
    {
        BinSearch(prefixD1, target.x, x_lower, x_upper, x_boundary, 0);
        index += x_lower * (GirdMax[1] - GridMin[1]);
    }

    if (target.y > (ld[1] / 2))
    {
        BinSearch(prefixD2, target.y, y_lower, y_upper, y_boundary, 1);
        index += y_upper;
    }
    else 
    {
        BinSearch(prefixD2, target.y, y_lower, y_upper, y_boundary, 0);
        index += y_lower;
    }

    if(target.z > (ld[2] / 2))
    {
        BinSearch(prefixD3, target.z, z_lower, z_upper, z_boundary, 1);
        index += z_upper * (GirdMax[1] - GridMin[1]) * (GirdMax[0] - GridMin[0]);
    }
    else
    {
        BinSearch(prefixD3, target.z, z_lower, z_upper, z_boundary, 0);
        index += z_lower * (GirdMax[1] - GridMin[1]) * (GirdMax[0] - GridMin[0]);
    }

    if (x_boundary && y_boundary && z_boundary)
    {
        // Assume y-x-z cycle
        target.s1 = SearchArray[index].s1;
        return 0;
    }

    // printf("target: %f %f %f, x_upper = %d, x_lower = %d, y_upper = %d, y_lower = %d, z_upper = %d, z_lower = %d\n", target.x, target.y, target.z, x_upper, x_lower, y_upper, y_lower, z_upper, z_lower);
    // Assume that the order of grid points follows y-x-z cycle
    {
        V c00, c01, c10, c11;
        V x_d, y_d, z_d;
        x_d = (target.x - prefixD1[x_lower])/(prefixD1[x_upper] - prefixD1[x_lower]);
        y_d = (target.y - prefixD2[y_lower])/(prefixD2[y_upper] - prefixD2[y_lower]);
        z_d = (target.z - prefixD3[z_lower])/(prefixD3[z_upper] - prefixD3[z_lower]);

        //(0, 0, 0) And (1, 0, 0)
        c00 = (1.0 - x_d) * SearchArray[YXZ(x_lower, y_lower, z_lower)].s1 + 
                     x_d * SearchArray[YXZ(x_upper, y_lower, z_lower)].s1;
        //(0, 0, 1) And (1, 0, 1)
        c01 = (1.0 - x_d) * SearchArray[YXZ(x_lower, y_lower, z_upper)].s1 + 
                     x_d * SearchArray[YXZ(x_upper, y_lower, z_upper)].s1;
        //(0, 1, 0) And (1, 1, 0)
        c10 = (1.0 - x_d) * SearchArray[YXZ(x_lower, y_upper, z_lower)].s1 + 
                     x_d * SearchArray[YXZ(x_upper, y_upper, z_lower)].s1;
        //(0, 1, 1) And (1, 1, 1)
        c11 = (1.0 - x_d) * SearchArray[YXZ(x_lower, y_upper, z_upper)].s1 + 
                     x_d * SearchArray[YXZ(x_upper, y_upper, z_upper)].s1;
        // c0 = (c00 * (1.0 - y_d) + c10 * y_d);
        // c1 = (c01 * (1.0 - y_d) + c11 * y_d);
        target.s1 = (c00 * (1.0 - y_d) + c10 * y_d) * (1.0 - z_d) + (c01 * (1.0 - y_d) + c11 * y_d) * z_d;
    }

    return 0;
}
